{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# (Original) Training at scale with the Vertex AI Training Service\n",
    "\n",
    "## Learning objectives\n",
    "  1. Learn how to organize your training code into a Python package\n",
    "  2. Train your model using cloud infrastructure via Google Cloud Vertex AI Training Service\n",
    "  3. (optional) Learn how to run your training package using Docker containers and push training Docker images on a Docker registry\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In this notebook we'll make the jump from training locally, to do training in the cloud. We'll take advantage of Google Cloud's [Vertex AI Training Service](https://cloud.google.com/vertex-ai/). \n",
    "\n",
    "Vertex AI Training Service is a managed service that allows the training and deployment of ML models without having to provision or maintain servers. The infrastructure is handled seamlessly by the managed service for us.\n",
    "\n",
    "Each learning objective will correspond to a __#TODO__ in the [student lab notebook](../labs/1_training_at_scale_vertex.ipynb) -- try to complete that notebook first before reviewing this solution notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensorflow==2.6.0\n",
      "Requirement already satisfied: tensorflow-gpu in /opt/conda/lib/python3.7/site-packages (2.10.0)\n",
      "Requirement already satisfied: gast<=0.4.0,>=0.2.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (0.4.0)\n",
      "Requirement already satisfied: libclang>=13.0.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (14.0.6)\n",
      "Requirement already satisfied: six>=1.12.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.15.0)\n",
      "Requirement already satisfied: flatbuffers>=2.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (22.9.24)\n",
      "Requirement already satisfied: numpy>=1.20 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.21.6)\n",
      "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (3.19.5)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (59.8.0)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (3.7.4.3)\n",
      "Requirement already satisfied: h5py>=2.9.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (3.1.0)\n",
      "Requirement already satisfied: tensorboard<2.11,>=2.10 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (2.10.1)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.48.1)\n",
      "Requirement already satisfied: keras-preprocessing>=1.1.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.1.2)\n",
      "Requirement already satisfied: tensorflow-estimator<2.11,>=2.10.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (2.10.0)\n",
      "Requirement already satisfied: opt-einsum>=2.3.2 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (3.3.0)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (0.2.0)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.12.1)\n",
      "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (0.27.0)\n",
      "Requirement already satisfied: absl-py>=1.0.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.2.0)\n",
      "Requirement already satisfied: keras<2.11,>=2.10.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (2.10.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.1.0)\n",
      "Requirement already satisfied: packaging in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (21.3)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorflow-gpu) (1.6.3)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /opt/conda/lib/python3.7/site-packages (from astunparse>=1.6.0->tensorflow-gpu) (0.37.1)\n",
      "Requirement already satisfied: cached-property in /opt/conda/lib/python3.7/site-packages (from h5py>=2.9.0->tensorflow-gpu) (1.5.2)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (1.8.1)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (2.2.2)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (1.35.0)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (2.28.1)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (0.6.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (0.4.6)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /opt/conda/lib/python3.7/site-packages (from tensorboard<2.11,>=2.10->tensorflow-gpu) (3.4.1)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.7/site-packages (from packaging->tensorflow-gpu) (3.0.9)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /opt/conda/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow-gpu) (4.9)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow-gpu) (0.2.7)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.7/site-packages (from google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow-gpu) (4.2.4)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /opt/conda/lib/python3.7/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.11,>=2.10->tensorflow-gpu) (1.3.1)\n",
      "Requirement already satisfied: importlib-metadata>=4.4 in /opt/conda/lib/python3.7/site-packages (from markdown>=2.6.8->tensorboard<2.11,>=2.10->tensorflow-gpu) (4.11.4)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /opt/conda/lib/python3.7/site-packages (from requests<3,>=2.21.0->tensorboard<2.11,>=2.10->tensorflow-gpu) (2.1.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from requests<3,>=2.21.0->tensorboard<2.11,>=2.10->tensorflow-gpu) (1.26.11)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.7/site-packages (from requests<3,>=2.21.0->tensorboard<2.11,>=2.10->tensorflow-gpu) (3.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.7/site-packages (from requests<3,>=2.21.0->tensorboard<2.11,>=2.10->tensorflow-gpu) (2022.9.14)\n",
      "Requirement already satisfied: MarkupSafe>=2.1.1 in /opt/conda/lib/python3.7/site-packages (from werkzeug>=1.0.1->tensorboard<2.11,>=2.10->tensorflow-gpu) (2.1.1)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.7/site-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.11,>=2.10->tensorflow-gpu) (3.8.1)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /opt/conda/lib/python3.7/site-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.11,>=2.10->tensorflow-gpu) (0.4.8)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.7/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.11,>=2.10->tensorflow-gpu) (3.2.1)\n"
     ]
    }
   ],
   "source": [
    "!pip freeze | grep tensorflow==2.6.0 || pip install tensorflow==2.6.0\n",
    "!pip install tensorflow-gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** **Restart** the kernel before proceeding further (On the Notebook menu, select Kernel > Restart Kernel > Restart).\n",
    "\n",
    "Specify your project name and bucket name in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google import api_core\n",
    "from google.cloud import bigquery"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change the following cell as necessary:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: PROJECT=qwiklabs-gcp-02-4a5dc70b94ac\n",
      "env: BUCKET=qwiklabs-gcp-02-4a5dc70b94ac\n",
      "env: REGION=us-central1\n",
      "env: OUTDIR=gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data\n",
      "env: TFVERSION=2.8\n"
     ]
    }
   ],
   "source": [
    "# Change below if necessary\n",
    "PROJECT = !gcloud config get-value project  # noqa: E999\n",
    "PROJECT = PROJECT[0]\n",
    "BUCKET = PROJECT\n",
    "REGION = \"us-central1\"\n",
    "\n",
    "OUTDIR = f\"gs://{BUCKET}/taxifare/data\"\n",
    "\n",
    "%env PROJECT=$PROJECT\n",
    "%env BUCKET=$BUCKET\n",
    "%env REGION=$REGION\n",
    "%env OUTDIR=$OUTDIR\n",
    "%env TFVERSION=2.8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirm below that the bucket is regional and its region equals to the specified region:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://qwiklabs-gcp-02-4a5dc70b94ac/ :\n",
      "\tLocation type:\t\t\tregion\n",
      "\tLocation constraint:\t\tUS-CENTRAL1\n",
      "us-central1\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "gcloud storage ls --full --buckets gs://$BUCKET | grep \"gs://\\|Location\"\n",    "echo $REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Updated property [core/project].\n",
      "Updated property [ai/region].\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "gcloud config set project $PROJECT\n",
    "gcloud config set ai/region $REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create BigQuery tables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have not already created a BigQuery dataset for our data, run the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset already exists\n"
     ]
    }
   ],
   "source": [
    "bq = bigquery.Client(project=PROJECT)\n",
    "dataset = bigquery.Dataset(bq.dataset(\"taxifare\"))\n",
    "\n",
    "try:\n",
    "    bq.create_dataset(dataset)\n",
    "    print(\"Dataset created\")\n",
    "except api_core.exceptions.Conflict:\n",
    "    print(\"Dataset already exists\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a table with 1 million examples.\n",
    "\n",
    "Note that the order of columns is exactly what was in our CSV files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Query complete after 0.01s: 100%|██████████| 3/3 [00:00<00:00, 1553.06query/s]                        \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: []"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%bigquery\n",
    "\n",
    "CREATE OR REPLACE TABLE taxifare.feateng_training_data AS\n",
    "\n",
    "SELECT\n",
    "    (tolls_amount + fare_amount) AS fare_amount,\n",
    "    pickup_datetime,\n",
    "    pickup_longitude AS pickuplon,\n",
    "    pickup_latitude AS pickuplat,\n",
    "    dropoff_longitude AS dropofflon,\n",
    "    dropoff_latitude AS dropofflat,\n",
    "    passenger_count*1.0 AS passengers,\n",
    "    'unused' AS key\n",
    "FROM `nyc-tlc.yellow.trips`\n",
    "WHERE ABS(MOD(FARM_FINGERPRINT(CAST(pickup_datetime AS STRING)), 1000)) = 1\n",
    "AND\n",
    "    trip_distance > 0\n",
    "    AND fare_amount >= 2.5\n",
    "    AND pickup_longitude > -78\n",
    "    AND pickup_longitude < -70\n",
    "    AND dropoff_longitude > -78\n",
    "    AND dropoff_longitude < -70\n",
    "    AND pickup_latitude > 37\n",
    "    AND pickup_latitude < 45\n",
    "    AND dropoff_latitude > 37\n",
    "    AND dropoff_latitude < 45\n",
    "    AND passenger_count > 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make the validation dataset be 1/10 the size of the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Query complete after 0.00s: 100%|██████████| 3/3 [00:00<00:00, 1647.62query/s]                        \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Empty DataFrame\n",
       "Columns: []\n",
       "Index: []"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%bigquery\n",
    "\n",
    "CREATE OR REPLACE TABLE taxifare.feateng_valid_data AS\n",
    "\n",
    "SELECT\n",
    "    (tolls_amount + fare_amount) AS fare_amount,\n",
    "    pickup_datetime,\n",
    "    pickup_longitude AS pickuplon,\n",
    "    pickup_latitude AS pickuplat,\n",
    "    dropoff_longitude AS dropofflon,\n",
    "    dropoff_latitude AS dropofflat,\n",
    "    passenger_count*1.0 AS passengers,\n",
    "    'unused' AS key\n",
    "FROM `nyc-tlc.yellow.trips`\n",
    "WHERE ABS(MOD(FARM_FINGERPRINT(CAST(pickup_datetime AS STRING)), 10000)) = 2\n",
    "AND\n",
    "    trip_distance > 0\n",
    "    AND fare_amount >= 2.5\n",
    "    AND pickup_longitude > -78\n",
    "    AND pickup_longitude < -70\n",
    "    AND dropoff_longitude > -78\n",
    "    AND dropoff_longitude < -70\n",
    "    AND pickup_latitude > 37\n",
    "    AND pickup_latitude < 45\n",
    "    AND dropoff_latitude > 37\n",
    "    AND dropoff_latitude < 45\n",
    "    AND passenger_count > 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export the tables as CSV files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deleting current contents of gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data\n",
      "Extracting training data to gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data\n",
      "Extracting validation data to gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data\n",
      "  88345235  2022-09-28T08:55:48Z  gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data/taxi-train-000000000000.csv\n",
      "   8725746  2022-09-28T08:55:59Z  gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data/taxi-valid-000000000000.csv\n",
      "TOTAL: 2 objects, 97070981 bytes (92.57 MiB)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Waiting on bqjob_r639d330412682390_00000183834eea4d_1 ... (23s) Current status: DONE   \n",
      "Waiting on bqjob_r45ae6cca1ff7873_00000183834f507b_1 ... (2s) Current status: DONE   \n"
     ]
    }
   ],
   "source": [
    "# TODO 1 and TODO 2\n",
    "%%bash\n",
    "\n",
    "echo \"Deleting current contents of $OUTDIR\"\n",
    "gcloud storage rm --recursive --continue-on-error $OUTDIR\n",    "\n",
    "echo \"Extracting training data to $OUTDIR\"\n",
    "bq --location=US extract \\\n",
    "   --destination_format CSV  \\\n",
    "   --field_delimiter \",\" --noprint_header \\\n",
    "   taxifare.feateng_training_data \\\n",
    "   $OUTDIR/taxi-train-*.csv\n",
    "\n",
    "echo \"Extracting validation data to $OUTDIR\"\n",
    "bq --location=US extract \\\n",
    "   --destination_format CSV  \\\n",
    "   --field_delimiter \",\" --noprint_header \\\n",
    "   taxifare.feateng_valid_data \\\n",
    "   $OUTDIR/taxi-valid-*.csv\n",
    "\n",
    "gcloud storage ls --long $OUTDIR"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Confirm that you have created both the training and validation datasets in Google Cloud Storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data/taxi-train-000000000000.csv\n",
      "gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/data/taxi-valid-000000000000.csv\n"
     ]
    }
   ],
   "source": [
    "!gcloud storage ls gs://$BUCKET/taxifare/data"   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "37.85,2009-06-25 22:20:08 UTC,-73.972125,40.758541,-73.764108,40.762433,2,unused\n",
      "40.33,2015-01-07 13:29:40 UTC,-73.866294860839844,40.771163940429688,-73.9787368774414,40.756568908691406,2,unused\n"
     ]
    }
   ],
   "source": [
    "!gcloud storage cat gs://$BUCKET/taxifare/data/taxi-train-000000000000.csv | head -2"   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.5,2014-12-07 00:57:00 UTC,-73.971717,40.762905,-73.970867,40.762348,2,unused\n",
      "2.5,2009-07-02 23:41:00 UTC,-73.973312,40.750067,-73.97313,40.750078,2,unused\n"
     ]
    }
   ],
   "source": [
    "!gcloud storage cat gs://$BUCKET/taxifare/data/taxi-valid-000000000000.csv | head -2"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Make code compatible with Vertex AI Training Service\n",
    "In order to make our code compatible with Vertex AI Training Service we need to make the following changes:\n",
    "\n",
    "1. Upload data to Google Cloud Storage \n",
    "2. Move code into a trainer Python package\n",
    "4. Submit training job with `gcloud` to train on Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Move code into a python package\n",
    "\n",
    "The first thing to do is to convert your training code snippets into a regular Python package. \n",
    "\n",
    "A Python package is simply a collection of one or more `.py` files along with an `__init__.py` file to identify the containing directory as a package. The `__init__.py` sometimes contains initialization code but for our purposes an empty file suffices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create the package directory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our package directory contains 3 files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__init__.py  model.py  task.py\n"
     ]
    }
   ],
   "source": [
    "ls ./taxifare/trainer/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Paste existing code into model.py\n",
    "\n",
    "A Python package requires our code to be in a .py file, as opposed to notebook cells. So, we simply copy and paste our existing code for the previous notebook into a single file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below, we write the contents of the cell into `model.py` packaging the model we \n",
    "developed in the previous labs so that we can deploy it to Vertex AI Training Service.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./taxifare/trainer/model.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./taxifare/trainer/model.py\n",
    "\"\"\"Data prep, train and evaluate DNN model.\"\"\"\n",
    "\n",
    "import logging\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import callbacks, models\n",
    "from tensorflow.keras.layers import (\n",
    "    Concatenate,\n",
    "    Dense,\n",
    "    Discretization,\n",
    "    Embedding,\n",
    "    Flatten,\n",
    "    Input,\n",
    "    Lambda,\n",
    ")\n",
    "from tensorflow.keras.layers.experimental.preprocessing import HashedCrossing\n",
    "\n",
    "logging.info(tf.version.VERSION)\n",
    "\n",
    "CSV_COLUMNS = [\n",
    "    \"fare_amount\",\n",
    "    \"pickup_datetime\",\n",
    "    \"pickup_longitude\",\n",
    "    \"pickup_latitude\",\n",
    "    \"dropoff_longitude\",\n",
    "    \"dropoff_latitude\",\n",
    "    \"passenger_count\",\n",
    "    \"key\",\n",
    "]\n",
    "\n",
    "LABEL_COLUMN = \"fare_amount\"\n",
    "DEFAULTS = [[0.0], [\"na\"], [0.0], [0.0], [0.0], [0.0], [0.0], [\"na\"]]\n",
    "UNWANTED_COLS = [\"pickup_datetime\", \"key\"]\n",
    "\n",
    "INPUT_COLS = [\n",
    "    c for c in CSV_COLUMNS if c != LABEL_COLUMN and c not in UNWANTED_COLS\n",
    "]\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    for unwanted_col in UNWANTED_COLS:\n",
    "        row_data.pop(unwanted_col)\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "    return row_data, label\n",
    "\n",
    "\n",
    "def load_dataset(pattern, batch_size, num_repeat):\n",
    "    dataset = tf.data.experimental.make_csv_dataset(\n",
    "        file_pattern=pattern,\n",
    "        batch_size=batch_size,\n",
    "        column_names=CSV_COLUMNS,\n",
    "        column_defaults=DEFAULTS,\n",
    "        num_epochs=num_repeat,\n",
    "        shuffle_buffer_size=1000000,\n",
    "    )\n",
    "    return dataset.map(features_and_labels)\n",
    "\n",
    "\n",
    "def create_train_dataset(pattern, batch_size):\n",
    "    dataset = load_dataset(pattern, batch_size, num_repeat=None)\n",
    "    return dataset.prefetch(1)\n",
    "\n",
    "\n",
    "def create_eval_dataset(pattern, batch_size):\n",
    "    dataset = load_dataset(pattern, batch_size, num_repeat=1)\n",
    "    return dataset.prefetch(1)\n",
    "\n",
    "\n",
    "def euclidean(params):\n",
    "    lon1, lat1, lon2, lat2 = params\n",
    "    londiff = lon2 - lon1\n",
    "    latdiff = lat2 - lat1\n",
    "    return tf.sqrt(londiff * londiff + latdiff * latdiff)\n",
    "\n",
    "\n",
    "def scale_longitude(lon_column):\n",
    "    return (lon_column + 78) / 8.0\n",
    "\n",
    "\n",
    "def scale_latitude(lat_column):\n",
    "    return (lat_column - 37) / 8.0\n",
    "\n",
    "\n",
    "def transform(inputs, nbuckets):\n",
    "    transformed = {}\n",
    "\n",
    "    # Scaling longitude from range [-70, -78] to [0, 1]\n",
    "    transformed[\"scaled_plon\"] = Lambda(scale_longitude, name=\"scale_plon\")(\n",
    "        inputs[\"pickup_longitude\"]\n",
    "    )\n",
    "    transformed[\"scaled_dlon\"] = Lambda(scale_longitude, name=\"scale_dlon\")(\n",
    "        inputs[\"dropoff_longitude\"]\n",
    "    )\n",
    "\n",
    "    # Scaling latitude from range [37, 45] to [0, 1]\n",
    "    transformed[\"scaled_plat\"] = Lambda(scale_latitude, name=\"scale_plat\")(\n",
    "        inputs[\"pickup_latitude\"]\n",
    "    )\n",
    "    transformed[\"scaled_dlat\"] = Lambda(scale_latitude, name=\"scale_dlat\")(\n",
    "        inputs[\"dropoff_latitude\"]\n",
    "    )\n",
    "\n",
    "    # Apply euclidean function\n",
    "    transformed[\"euclidean_distance\"] = Lambda(euclidean, name=\"euclidean\")(\n",
    "        [\n",
    "            inputs[\"pickup_longitude\"],\n",
    "            inputs[\"pickup_latitude\"],\n",
    "            inputs[\"dropoff_longitude\"],\n",
    "            inputs[\"dropoff_latitude\"],\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    latbuckets = np.linspace(start=0.0, stop=1.0, num=nbuckets).tolist()\n",
    "    lonbuckets = np.linspace(start=0.0, stop=1.0, num=nbuckets).tolist()\n",
    "\n",
    "    # Bucketization with Discretization layer\n",
    "    plon = Discretization(lonbuckets, name=\"plon_bkt\")(\n",
    "        transformed[\"scaled_plon\"]\n",
    "    )\n",
    "    plat = Discretization(latbuckets, name=\"plat_bkt\")(\n",
    "        transformed[\"scaled_plat\"]\n",
    "    )\n",
    "    dlon = Discretization(lonbuckets, name=\"dlon_bkt\")(\n",
    "        transformed[\"scaled_dlon\"]\n",
    "    )\n",
    "    dlat = Discretization(latbuckets, name=\"dlat_bkt\")(\n",
    "        transformed[\"scaled_dlat\"]\n",
    "    )\n",
    "\n",
    "    # Feature Cross with HashedCrossing layer\n",
    "    p_fc = HashedCrossing(num_bins=nbuckets * nbuckets, name=\"p_fc\")(\n",
    "        (plon, plat)\n",
    "    )\n",
    "    d_fc = HashedCrossing(num_bins=nbuckets * nbuckets, name=\"d_fc\")(\n",
    "        (dlon, dlat)\n",
    "    )\n",
    "    pd_fc = HashedCrossing(num_bins=nbuckets**4, name=\"pd_fc\")((p_fc, d_fc))\n",
    "\n",
    "    # Embedding with Embedding layer\n",
    "    transformed[\"pd_embed\"] = Flatten()(\n",
    "        Embedding(input_dim=nbuckets**4, output_dim=10, name=\"pd_embed\")(\n",
    "            pd_fc\n",
    "        )\n",
    "    )\n",
    "\n",
    "    transformed[\"passenger_count\"] = inputs[\"passenger_count\"]\n",
    "\n",
    "    return transformed\n",
    "\n",
    "\n",
    "def rmse(y_true, y_pred):\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))\n",
    "\n",
    "\n",
    "def build_dnn_model(nbuckets, nnsize, lr):\n",
    "    inputs = {\n",
    "        colname: Input(name=colname, shape=(1,), dtype=\"float32\")\n",
    "        for colname in INPUT_COLS\n",
    "    }\n",
    "\n",
    "    # transforms\n",
    "    transformed = transform(inputs, nbuckets)\n",
    "    dnn_inputs = Concatenate()(transformed.values())\n",
    "\n",
    "    x = dnn_inputs\n",
    "    for layer, nodes in enumerate(nnsize):\n",
    "        x = Dense(nodes, activation=\"relu\", name=f\"h{layer}\")(x)\n",
    "    output = Dense(1, name=\"fare\")(x)\n",
    "\n",
    "    model = models.Model(inputs, output)\n",
    "    # TODO 1a\n",
    "    lr_optimizer = tf.keras.optimizers.Adam(learning_rate=lr)\n",
    "    model.compile(optimizer=lr_optimizer, loss=\"mse\", metrics=[rmse, \"mse\"])\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def train_and_evaluate(hparams):\n",
    "    # TODO 1b\n",
    "    batch_size = hparams[\"batch_size\"]\n",
    "    nbuckets = hparams[\"nbuckets\"]\n",
    "    lr = hparams[\"lr\"]\n",
    "    nnsize = [int(s) for s in hparams[\"nnsize\"].split()]\n",
    "    eval_data_path = hparams[\"eval_data_path\"]\n",
    "    num_evals = hparams[\"num_evals\"]\n",
    "    num_examples_to_train_on = hparams[\"num_examples_to_train_on\"]\n",
    "    output_dir = hparams[\"output_dir\"]\n",
    "    train_data_path = hparams[\"train_data_path\"]\n",
    "\n",
    "    model_export_path = os.path.join(output_dir, \"savedmodel\")\n",
    "    checkpoint_path = os.path.join(output_dir, \"checkpoints\")\n",
    "    tensorboard_path = os.path.join(output_dir, \"tensorboard\")\n",
    "\n",
    "    if tf.io.gfile.exists(output_dir):\n",
    "        tf.io.gfile.rmtree(output_dir)\n",
    "\n",
    "    model = build_dnn_model(nbuckets, nnsize, lr)\n",
    "    logging.info(model.summary())\n",
    "\n",
    "    trainds = create_train_dataset(train_data_path, batch_size)\n",
    "    evalds = create_eval_dataset(eval_data_path, batch_size)\n",
    "\n",
    "    steps_per_epoch = num_examples_to_train_on // (batch_size * num_evals)\n",
    "\n",
    "    checkpoint_cb = callbacks.ModelCheckpoint(\n",
    "        checkpoint_path, save_weights_only=True, verbose=1\n",
    "    )\n",
    "    tensorboard_cb = callbacks.TensorBoard(tensorboard_path, histogram_freq=1)\n",
    "\n",
    "    history = model.fit(\n",
    "        trainds,\n",
    "        validation_data=evalds,\n",
    "        epochs=num_evals,\n",
    "        steps_per_epoch=max(1, steps_per_epoch),\n",
    "        verbose=2,  # 0=silent, 1=progress bar, 2=one line per epoch\n",
    "        callbacks=[checkpoint_cb, tensorboard_cb],\n",
    "    )\n",
    "\n",
    "    # Exporting the model with default serving function.\n",
    "    model.save(model_export_path)\n",
    "    return history\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modify code to read data from and write checkpoint files to GCS \n",
    "\n",
    "If you look closely above, you'll notice a new function, `train_and_evaluate` that wraps the code that actually trains the model. This allows us to parametrize the training by passing a dictionary of parameters to this function (e.g, `batch_size`, `num_examples_to_train_on`, `train_data_path` etc.)\n",
    "\n",
    "This is useful because the output directory, data paths and number of train steps will be different depending on whether we're training locally or in the cloud. Parametrizing allows us to use the same code for both.\n",
    "\n",
    "We specify these parameters at run time via the command line. Which means we need to add code to parse command line parameters and invoke `train_and_evaluate()` with those params. This is the job of the `task.py` file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting taxifare/trainer/task.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile taxifare/trainer/task.py\n",
    "\"\"\"Argument definitions for model training code in `trainer.model`.\"\"\"\n",
    "\n",
    "import argparse\n",
    "\n",
    "from trainer import model\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\n",
    "        \"--batch_size\",\n",
    "        help=\"Batch size for training steps\",\n",
    "        type=int,\n",
    "        default=32,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_data_path\",\n",
    "        help=\"GCS location pattern of eval files\",\n",
    "        required=True,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nnsize\",\n",
    "        help=\"Hidden layer sizes (provide space-separated sizes)\",\n",
    "        default=\"32 8\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nbuckets\",\n",
    "        help=\"Number of buckets to divide lat and lon with\",\n",
    "        type=int,\n",
    "        default=10,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--lr\", help=\"learning rate for optimizer\", type=float, default=0.001\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_evals\",\n",
    "        help=\"Number of times to evaluate model on eval data training.\",\n",
    "        type=int,\n",
    "        default=5,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_examples_to_train_on\",\n",
    "        help=\"Number of examples to train on.\",\n",
    "        type=int,\n",
    "        default=100,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--output_dir\",\n",
    "        help=\"GCS location to write checkpoints and export models\",\n",
    "        required=True,\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--train_data_path\",\n",
    "        help=\"GCS location pattern of train files containing eval URLs\",\n",
    "        required=True,\n",
    "    )\n",
    "    args = parser.parse_args()\n",
    "    hparams = args.__dict__\n",
    "\n",
    "    model.train_and_evaluate(hparams)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run trainer module package locally\n",
    "\n",
    "Now we can test our training code locally as follows using the local test data. We'll run a very small training job over a single file with a small batch size and one eval step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " pickup_longitude (InputLayer)  [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " dropoff_longitude (InputLayer)  [(None, 1)]         0           []                               \n",
      "                                                                                                  \n",
      " pickup_latitude (InputLayer)   [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " dropoff_latitude (InputLayer)  [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " scale_plon (Lambda)            (None, 1)            0           ['pickup_longitude[0][0]']       \n",
      "                                                                                                  \n",
      " scale_dlon (Lambda)            (None, 1)            0           ['dropoff_longitude[0][0]']      \n",
      "                                                                                                  \n",
      " scale_plat (Lambda)            (None, 1)            0           ['pickup_latitude[0][0]']        \n",
      "                                                                                                  \n",
      " scale_dlat (Lambda)            (None, 1)            0           ['dropoff_latitude[0][0]']       \n",
      "                                                                                                  \n",
      " plon_bkt (Discretization)      (None, 1)            0           ['scale_plon[0][0]']             \n",
      "                                                                                                  \n",
      " plat_bkt (Discretization)      (None, 1)            0           ['scale_plat[0][0]']             \n",
      "                                                                                                  \n",
      " dlon_bkt (Discretization)      (None, 1)            0           ['scale_dlon[0][0]']             \n",
      "                                                                                                  \n",
      " dlat_bkt (Discretization)      (None, 1)            0           ['scale_dlat[0][0]']             \n",
      "                                                                                                  \n",
      " p_fc (HashedCrossing)          (None, 1)            0           ['plon_bkt[0][0]',               \n",
      "                                                                  'plat_bkt[0][0]']               \n",
      "                                                                                                  \n",
      " d_fc (HashedCrossing)          (None, 1)            0           ['dlon_bkt[0][0]',               \n",
      "                                                                  'dlat_bkt[0][0]']               \n",
      "                                                                                                  \n",
      " pd_fc (HashedCrossing)         (None, 1)            0           ['p_fc[0][0]',                   \n",
      "                                                                  'd_fc[0][0]']                   \n",
      "                                                                                                  \n",
      " pd_embed (Embedding)           (None, 1, 10)        100000      ['pd_fc[0][0]']                  \n",
      "                                                                                                  \n",
      " euclidean (Lambda)             (None, 1)            0           ['pickup_longitude[0][0]',       \n",
      "                                                                  'pickup_latitude[0][0]',        \n",
      "                                                                  'dropoff_longitude[0][0]',      \n",
      "                                                                  'dropoff_latitude[0][0]']       \n",
      "                                                                                                  \n",
      " flatten (Flatten)              (None, 10)           0           ['pd_embed[0][0]']               \n",
      "                                                                                                  \n",
      " passenger_count (InputLayer)   [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " concatenate (Concatenate)      (None, 16)           0           ['scale_plon[0][0]',             \n",
      "                                                                  'scale_dlon[0][0]',             \n",
      "                                                                  'scale_plat[0][0]',             \n",
      "                                                                  'scale_dlat[0][0]',             \n",
      "                                                                  'euclidean[0][0]',              \n",
      "                                                                  'flatten[0][0]',                \n",
      "                                                                  'passenger_count[0][0]']        \n",
      "                                                                                                  \n",
      " h0 (Dense)                     (None, 32)           544         ['concatenate[0][0]']            \n",
      "                                                                                                  \n",
      " h1 (Dense)                     (None, 8)            264         ['h0[0][0]']                     \n",
      "                                                                                                  \n",
      " fare (Dense)                   (None, 1)            9           ['h1[0][0]']                     \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 100,817\n",
      "Trainable params: 100,817\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n",
      "\n",
      "Epoch 1: saving model to ./taxifare-model/checkpoints\n",
      "20/20 - 15s - loss: 82.1748 - rmse: 8.0489 - mse: 82.1748 - val_loss: 86.1465 - val_rmse: 7.7346 - val_mse: 86.1465 - 15s/epoch - 743ms/step\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-09-28 08:56:56.548260: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-09-28 08:56:56.757628: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2022-09-28 08:56:57.677068: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64\n",
      "2022-09-28 08:56:57.677173: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64\n",
      "2022-09-28 08:56:57.677185: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
      "2022-09-28 08:56:58.930262: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64\n",
      "2022-09-28 08:56:58.930309: W tensorflow/stream_executor/cuda/cuda_driver.cc:263] failed call to cuInit: UNKNOWN ERROR (303)\n",
      "2022-09-28 08:56:58.930334: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (tensorflow-2-6-20220928-140445): /proc/driver/nvidia/version does not exist\n",
      "2022-09-28 08:56:58.930581: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "# TODO 3\n",
    "\n",
    "%%bash\n",
    "\n",
    "EVAL_DATA_PATH=../data/taxi-traffic-valid*\n",
    "TRAIN_DATA_PATH=../data/taxi-traffic-train*\n",
    "OUTPUT_DIR=./taxifare-model\n",
    "\n",
    "test ${OUTPUT_DIR} && rm -rf ${OUTPUT_DIR}\n",
    "export PYTHONPATH=${PYTHONPATH}:${PWD}/taxifare\n",
    "    \n",
    "# Run the trainer module package locally with 1 eval\n",
    "\n",
    "python3 -m trainer.task \\\n",
    "--eval_data_path $EVAL_DATA_PATH \\\n",
    "--output_dir $OUTPUT_DIR \\\n",
    "--train_data_path $TRAIN_DATA_PATH \\\n",
    "--batch_size 5 \\\n",
    "--num_examples_to_train_on 100 \\\n",
    "--num_evals 1 \\\n",
    "--nbuckets 10 \\\n",
    "--lr 0.001 \\\n",
    "--nnsize \"32 8\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run your training package on Vertex AI using a pre-built container\n",
    "\n",
    "Once the code works in standalone mode locally, you can run it on the Cloud using Vertex AI and use pre-built containers. First, we need to package our code as a source distribution. For this, we can use `setuptools`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting taxifare/setup.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile taxifare/setup.py\n",
    "\"\"\"Using `setuptools` to create a source distribution.\"\"\"\n",
    "\n",
    "from setuptools import find_packages, setup\n",
    "\n",
    "setup(\n",
    "    name=\"taxifare_trainer\",\n",
    "    version=\"0.1\",\n",
    "    packages=find_packages(),\n",
    "    include_package_data=True,\n",
    "    description=\"Taxifare model training application.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "running sdist\n",
      "running egg_info\n",
      "creating taxifare_trainer.egg-info\n",
      "writing taxifare_trainer.egg-info/PKG-INFO\n",
      "writing dependency_links to taxifare_trainer.egg-info/dependency_links.txt\n",
      "writing top-level names to taxifare_trainer.egg-info/top_level.txt\n",
      "writing manifest file 'taxifare_trainer.egg-info/SOURCES.txt'\n",
      "reading manifest file 'taxifare_trainer.egg-info/SOURCES.txt'\n",
      "writing manifest file 'taxifare_trainer.egg-info/SOURCES.txt'\n",
      "running check\n",
      "creating taxifare_trainer-0.1\n",
      "creating taxifare_trainer-0.1/taxifare_trainer.egg-info\n",
      "creating taxifare_trainer-0.1/trainer\n",
      "copying files to taxifare_trainer-0.1...\n",
      "copying setup.py -> taxifare_trainer-0.1\n",
      "copying taxifare_trainer.egg-info/PKG-INFO -> taxifare_trainer-0.1/taxifare_trainer.egg-info\n",
      "copying taxifare_trainer.egg-info/SOURCES.txt -> taxifare_trainer-0.1/taxifare_trainer.egg-info\n",
      "copying taxifare_trainer.egg-info/dependency_links.txt -> taxifare_trainer-0.1/taxifare_trainer.egg-info\n",
      "copying taxifare_trainer.egg-info/top_level.txt -> taxifare_trainer-0.1/taxifare_trainer.egg-info\n",
      "copying trainer/__init__.py -> taxifare_trainer-0.1/trainer\n",
      "copying trainer/model.py -> taxifare_trainer-0.1/trainer\n",
      "copying trainer/task.py -> taxifare_trainer-0.1/trainer\n",
      "Writing taxifare_trainer-0.1/setup.cfg\n",
      "creating dist\n",
      "Creating tar archive\n",
      "removing 'taxifare_trainer-0.1' (and everything under it)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "warning: sdist: standard file not found: should have one of README, README.rst, README.txt, README.md\n",
      "\n",
      "warning: check: missing required meta-data: url\n",
      "\n",
      "warning: check: missing meta-data: either (author and author_email) or (maintainer and maintainer_email) must be supplied\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "cd taxifare\n",
    "python ./setup.py sdist --formats=gztar\n",
    "cd .."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will store our package in the Cloud Storage bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Copying file://taxifare/dist/taxifare_trainer-0.1.tar.gz [Content-Type=application/x-tar]...\n",
      "/ [1 files][  3.2 KiB/  3.2 KiB]                                                \n",
      "Operation completed over 1 objects/3.2 KiB.                                      \n"
     ]
    }
   ],
   "source": [
    "# TODO 4\n",
    "# Store the package in the Cloud Storage bucket\n",
    "\n",
    "%%bash\n",
    "gcloud storage cp taxifare/dist/taxifare_trainer-0.1.tar.gz gs://${BUCKET}/taxifare/"   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Submit Custom Job using the `gcloud` CLI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To submit this source distribution the Cloud we use [`gcloud ai custom-jobs create`](https://cloud.google.com/sdk/gcloud/reference/ai/custom-jobs/create) and simply specify some additional parameters for Vertex AI Training service:\n",
    "- job_name: A unique identifier for the Cloud job. We usually append system time to ensure uniqueness\n",
    "- region: Cloud region to train in. See [here](https://cloud.google.com/vertex-ai/docs/general/locations) for supported Vertex AI Custom model training regions\n",
    "\n",
    "The arguments within `--args` are sent to our `task.py`.\n",
    "\n",
    "Because this is on the entire dataset, it will take a while. You can monitor the job from the GCP console in the Vertex AI Training section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://qwiklabs-gcp-02-4a5dc70b94ac/taxifare/trained_model_20220928_090152 us-central1 taxifare_20220928_090152\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using endpoint [https://us-central1-aiplatform.googleapis.com/]\n",
      "CustomJob [projects/429871156727/locations/us-central1/customJobs/1740844390751928320] is submitted successfully.\n",
      "\n",
      "Your job is still active. You may view the status of your job with the command\n",
      "\n",
      "  $ gcloud ai custom-jobs describe projects/429871156727/locations/us-central1/customJobs/1740844390751928320\n",
      "\n",
      "or continue streaming the logs with the command\n",
      "\n",
      "  $ gcloud ai custom-jobs stream-logs projects/429871156727/locations/us-central1/customJobs/1740844390751928320\n"
     ]
    }
   ],
   "source": [
    "# TODO 5\n",
    "%%bash\n",
    "\n",
    "# Output directory and jobID\n",
    "TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)\n",
    "OUTDIR=gs://${BUCKET}/taxifare/trained_model_$TIMESTAMP\n",
    "JOB_NAME=taxifare_$TIMESTAMP\n",
    "echo ${OUTDIR} ${REGION} ${JOB_NAME}\n",
    "\n",
    "PYTHON_PACKAGE_URIS=gs://${BUCKET}/taxifare/taxifare_trainer-0.1.tar.gz\n",
    "MACHINE_TYPE=n1-standard-4\n",
    "REPLICA_COUNT=1\n",
    "PYTHON_PACKAGE_EXECUTOR_IMAGE_URI=\"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\"\n",
    "PYTHON_MODULE=trainer.task\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=50\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=5000\n",
    "NUM_EVALS=100\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "# GCS paths\n",
    "GCS_PROJECT_PATH=gs://$BUCKET/taxifare\n",
    "DATA_PATH=$GCS_PROJECT_PATH/data\n",
    "TRAIN_DATA_PATH=$DATA_PATH/taxi-train*\n",
    "EVAL_DATA_PATH=$DATA_PATH/taxi-valid*\n",
    "\n",
    "WORKER_POOL_SPEC=\"machine-type=$MACHINE_TYPE,\\\n",
    "replica-count=$REPLICA_COUNT,\\\n",
    "executor-image-uri=$PYTHON_PACKAGE_EXECUTOR_IMAGE_URI,\\\n",
    "python-module=$PYTHON_MODULE\"\n",
    "\n",
    "ARGS=\"--eval_data_path=$EVAL_DATA_PATH,\\\n",
    "--output_dir=$OUTDIR,\\\n",
    "--train_data_path=$TRAIN_DATA_PATH,\\\n",
    "--batch_size=$BATCH_SIZE,\\\n",
    "--num_examples_to_train_on=$NUM_EXAMPLES_TO_TRAIN_ON,\\\n",
    "--num_evals=$NUM_EVALS,\\\n",
    "--nbuckets=$NBUCKETS,\\\n",
    "--lr=$LR,\\\n",
    "--nnsize=$NNSIZE\"\n",
    "\n",
    "# Create a custom job\n",
    "\n",
    "gcloud ai custom-jobs create \\\n",
    "  --region=${REGION} \\\n",
    "  --display-name=$JOB_NAME \\\n",
    "  --python-package-uris=$PYTHON_PACKAGE_URIS \\\n",
    "  --worker-pool-spec=$WORKER_POOL_SPEC \\\n",
    "  --args=\"$ARGS\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Submit Custom Job using the Vertex AI Python SDK"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `gcloud` CLI is just one of multiple ways to interact with Vertex AI, which also include the Console GUI, directly calling the REST APIs (e.g. using `curl`), and the most flexible interface being the Vertex AI SDK available in multiple languages.\n",
    "\n",
    "Below, we use the Vertex AI Python SDK to accomplish the same Pre-Built Container training as above -- getting familiar with the SDK will come in handy later when we use advanced features such as hyperparameter tuning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from datetime import datetime\n",
    "from google.cloud import aiplatform\n",
    "\n",
    "\n",
    "# Output directory and jobID\n",
    "timestamp_str=datetime.strftime(datetime.now(), '%y%m%d_%H%M%S')\n",
    "OUTDIR=f\"gs://{BUCKET}/taxifare/trained_model_{timestamp_str}\"\n",
    "JOB_NAME=f\"taxifare_{timestamp_str}\"\n",
    "print(OUTDIR, REGION, JOB_NAME)\n",
    "\n",
    "\n",
    "# Vertex AI machines to use for training\n",
    "PYTHON_PACKAGE_URIS=f\"gs://{BUCKET}/taxifare/taxifare_trainer-0.1.tar.gz\"\n",
    "MACHINE_TYPE=\"n1-standard-4\"\n",
    "REPLICA_COUNT=1\n",
    "PYTHON_PACKAGE_EXECUTOR_IMAGE_URI=\"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\"\n",
    "PYTHON_MODULE=\"trainer.task\"\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=50\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=5000\n",
    "NUM_EVALS=100\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "\n",
    "# GCS paths.\n",
    "GCS_PROJECT_PATH=f\"gs://{BUCKET}/taxifare\"\n",
    "DATA_PATH=f\"{GCS_PROJECT_PATH}/data\"\n",
    "TRAIN_DATA_PATH=f\"{DATA_PATH}/taxi-train*\"\n",
    "EVAL_DATA_PATH=f\"{DATA_PATH}/taxi-valid*\"\n",
    "\n",
    "def create_custom_job_sample(\n",
    "    project: str,\n",
    "    display_name: str,\n",
    "    executor_image_uri: str,\n",
    "    package_uri: str,\n",
    "    python_module: str,\n",
    "    location: str = REGION,\n",
    "    api_endpoint: str = f\"{REGION}-aiplatform.googleapis.com\",\n",
    "):\n",
    "    # The AI Platform services require regional API endpoints.\n",
    "    client_options = {\"api_endpoint\": api_endpoint}\n",
    "    # Initialize client that will be used to create and send requests.\n",
    "    # This client only needs to be created once, and can be reused for multiple requests.\n",
    "    client = aiplatform.gapic.JobServiceClient(client_options=client_options)\n",
    "    custom_job = {\n",
    "        \"display_name\": display_name,\n",
    "        \"job_spec\": {\n",
    "            \"worker_pool_specs\": [\n",
    "                {\n",
    "                    \"machine_spec\": {\n",
    "                        \"machine_type\": MACHINE_TYPE,\n",
    "                    },\n",
    "                    \"replica_count\": REPLICA_COUNT,\n",
    "                    \"python_package_spec\": {\n",
    "                        \"executor_image_uri\": executor_image_uri,\n",
    "                        \"package_uris\": [package_uri],\n",
    "                        \"python_module\": python_module,\n",
    "                        \"args\": [\n",
    "                            f\"--eval_data_path={EVAL_DATA_PATH}\",\n",
    "                            f\"--output_dir={OUTDIR}\",\n",
    "                            f\"--train_data_path={TRAIN_DATA_PATH}\",\n",
    "                            f\"--batch_size={BATCH_SIZE}\",\n",
    "                            f\"--num_examples_to_train_on={NUM_EXAMPLES_TO_TRAIN_ON}\",\n",
    "                            f\"--num_evals={NUM_EVALS}\",\n",
    "                            f\"--nbuckets={NBUCKETS}\",\n",
    "                            f\"--lr={LR}\",\n",
    "                            f\"--nnsize={NNSIZE}\"\n",
    "                        ],\n",
    "                    },\n",
    "                }\n",
    "            ],\n",
    "        },\n",
    "    }\n",
    "    parent = f\"projects/{project}/locations/{location}\"\n",
    "    response = client.create_custom_job(parent=parent, custom_job=custom_job)\n",
    "    print(\"response:\", response)\n",
    "\n",
    "create_custom_job_sample(\n",
    "    project=PROJECT,\n",
    "    display_name=JOB_NAME,\n",
    "    executor_image_uri=PYTHON_PACKAGE_EXECUTOR_IMAGE_URI,\n",
    "    package_uri=PYTHON_PACKAGE_URIS,\n",
    "    python_module=PYTHON_MODULE)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run your training package using a custom container\n",
    "\n",
    "Vertex AI Training also supports training in custom containers, allowing users to bring their own Docker containers with any pre-installed ML framework or algorithm to run on Vertex AI Training. \n",
    "\n",
    "In this last section, we'll see how to submit a Cloud training job using a customized Docker image. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Containerizing our `./taxifare/trainer` package involves 3 steps:\n",
    "\n",
    "* Writing a Dockerfile in `./taxifare`\n",
    "* Building the Docker image\n",
    "* Pushing it to the Google Cloud container registry in our GCP project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `Dockerfile` specifies\n",
    "1. How the container needs to be provisioned so that all the dependencies in our code are satisfied\n",
    "2. Where to copy our trainer Package in the container\n",
    "3. What command to run when the container is ran (the `ENTRYPOINT` line)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./taxifare/Dockerfile\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./taxifare/Dockerfile\n",
    "FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\n",
    "# TODO 3\n",
    "\n",
    "COPY . /code\n",
    "\n",
    "WORKDIR /code\n",
    "\n",
    "ENTRYPOINT [\"python3\", \"-m\", \"trainer.task\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%%bash \n",
    "\n",
    "PROJECT_DIR=$(cd ./taxifare && pwd)\n",
    "IMAGE_NAME=taxifare_training_container\n",
    "DOCKERFILE=$PROJECT_DIR/Dockerfile\n",
    "IMAGE_URI=gcr.io/$PROJECT/$IMAGE_NAME\n",
    "\n",
    "docker build $PROJECT_DIR -f $DOCKERFILE -t $IMAGE_URI\n",
    "\n",
    "docker push $IMAGE_URI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Remark:** If you prefer to build the container image from the command line, we have written a script for that `./taxifare/scripts/build.sh`. This script reads its configuration from the file `./taxifare/scripts/env.sh`. You can configure these arguments the way you want in that file. You can also simply type `make build` from within `./taxifare` to build the image (which will invoke the build script). Similarly, we wrote the script `./taxifare/scripts/push.sh` to push the Docker image, which you can also trigger by typing `make push` from within `./taxifare`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train using a custom container on Vertex AI\n",
    "\n",
    "TODO: To submit to the Cloud we use [`gcloud ai custom-jobs create`](https://cloud.google.com/sdk/gcloud/reference/ai/custom-jobs/create) and simply specify some additional parameters for Vertex AI Training Service:\n",
    "- job_name: A unique identifier for the Cloud job. We usually append system time to ensure uniqueness\n",
    "- image-uri: The uri of the Docker image we pushed in the Google Cloud registry\n",
    "- region: Cloud region to train in. See [here](https://cloud.google.com/vertex-ai/docs/general/locations) for supported Vertex AI Training Service regions\n",
    "\n",
    "The arguments within `--args` are sent to our `task.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can track your job and view logs using [cloud console](https://console.cloud.google.com/mlengine/jobs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "flake8-noqa-cell"
    ]
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "# Output directory and jobID\n",
    "TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)\n",
    "OUTDIR=gs://${BUCKET}/taxifare/trained_model_$TIMESTAMP\n",
    "JOB_NAME=taxifare_container_$TIMESTAMP\n",
    "echo ${OUTDIR} ${REGION} ${JOB_NAME}\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=50\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=5000\n",
    "NUM_EVALS=100\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "# Vertex AI machines to use for training\n",
    "MACHINE_TYPE=n1-standard-4\n",
    "REPLICA_COUNT=1\n",
    "\n",
    "# GCS paths.\n",
    "GCS_PROJECT_PATH=gs://$BUCKET/taxifare\n",
    "DATA_PATH=$GCS_PROJECT_PATH/data\n",
    "TRAIN_DATA_PATH=$DATA_PATH/taxi-train*\n",
    "EVAL_DATA_PATH=$DATA_PATH/taxi-valid*\n",
    "\n",
    "IMAGE_NAME=taxifare_training_container\n",
    "IMAGE_URI=gcr.io/$PROJECT/$IMAGE_NAME\n",
    "\n",
    "WORKER_POOL_SPEC=\"machine-type=$MACHINE_TYPE,\\\n",
    "replica-count=$REPLICA_COUNT,\\\n",
    "container-image-uri=$IMAGE_URI\"\n",
    "\n",
    "ARGS=\"--eval_data_path=$EVAL_DATA_PATH,\\\n",
    "--output_dir=$OUTDIR,\\\n",
    "--train_data_path=$TRAIN_DATA_PATH,\\\n",
    "--batch_size=$BATCH_SIZE,\\\n",
    "--num_examples_to_train_on=$NUM_EXAMPLES_TO_TRAIN_ON,\\\n",
    "--num_evals=$NUM_EVALS,\\\n",
    "--nbuckets=$NBUCKETS,\\\n",
    "--lr=$LR,\\\n",
    "--nnsize=$NNSIZE\"\n",
    "\n",
    "gcloud ai custom-jobs create \\\n",
    "  --region=$REGION \\\n",
    "  --display-name=$JOB_NAME \\\n",
    "  --worker-pool-spec=$WORKER_POOL_SPEC \\\n",
    "  --args=\"$ARGS\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As seen with the Pre-built Container, we can use the Vertex AI Python SDK also here as an alternative to the `gcloud` CLI as below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "from datetime import datetime\n",
    "from google.cloud import aiplatform\n",
    "\n",
    "\n",
    "# Output directory and jobID\n",
    "timestamp_str=datetime.strftime(datetime.now(), '%Y%m%d_%H%M%S')\n",
    "OUTDIR=f\"gs://{BUCKET}/taxifare/trained_model_{timestamp_str}\"\n",
    "JOB_NAME=f\"taxifare_{timestamp_str}\"\n",
    "print(OUTDIR, REGION, JOB_NAME)\n",
    "\n",
    "\n",
    "# Vertex AI machines to use for training\n",
    "MACHINE_TYPE=\"n1-standard-4\"\n",
    "REPLICA_COUNT=1\n",
    "IMAGE_NAME=\"taxifare_training_container\"\n",
    "IMAGE_URI=f\"gcr.io/{PROJECT}/{IMAGE_NAME}\"\n",
    "\n",
    "# Model and training hyperparameters\n",
    "BATCH_SIZE=50\n",
    "NUM_EXAMPLES_TO_TRAIN_ON=5000\n",
    "NUM_EVALS=100\n",
    "NBUCKETS=10\n",
    "LR=0.001\n",
    "NNSIZE=\"32 8\"\n",
    "\n",
    "\n",
    "# GCS paths.\n",
    "GCS_PROJECT_PATH=f\"gs://{BUCKET}/taxifare\"\n",
    "DATA_PATH=f\"{GCS_PROJECT_PATH}/data\"\n",
    "TRAIN_DATA_PATH=f\"{DATA_PATH}/taxi-train*\"\n",
    "EVAL_DATA_PATH=f\"{DATA_PATH}/taxi-valid*\"\n",
    "\n",
    "# custom container\n",
    "IMAGE_NAME=\"taxifare_training_container\"\n",
    "IMAGE_URI=f\"gcr.io/{PROJECT}/{IMAGE_NAME}\"\n",
    "\n",
    "def create_custom_job_sample(\n",
    "    project: str,\n",
    "    display_name: str,\n",
    "    image_uri: str,\n",
    "    location: str = REGION,\n",
    "    api_endpoint: str = f\"{REGION}-aiplatform.googleapis.com\",\n",
    "):\n",
    "    # The AI Platform services require regional API endpoints.\n",
    "    client_options = {\"api_endpoint\": api_endpoint}\n",
    "    # Initialize client that will be used to create and send requests.\n",
    "    # This client only needs to be created once, and can be reused for multiple requests.\n",
    "    client = aiplatform.gapic.JobServiceClient(client_options=client_options)\n",
    "    custom_job = {\n",
    "        \"display_name\": display_name,\n",
    "        \"job_spec\": {\n",
    "            \"worker_pool_specs\": [\n",
    "                {\n",
    "                    \"machine_spec\": {\n",
    "                        \"machine_type\": \"n1-standard-4\",\n",
    "                    },\n",
    "                    \"replica_count\": 1,\n",
    "                    \"container_spec\": {\n",
    "                        \"image_uri\": image_uri,\n",
    "                        \"args\": [\n",
    "                            f\"--eval_data_path={EVAL_DATA_PATH}\",\n",
    "                            f\"--output_dir={OUTDIR}\",\n",
    "                            f\"--train_data_path={TRAIN_DATA_PATH}\",\n",
    "                            f\"--batch_size={BATCH_SIZE}\",\n",
    "                            f\"--num_examples_to_train_on={NUM_EXAMPLES_TO_TRAIN_ON}\",\n",
    "                            f\"--num_evals={NUM_EVALS}\",\n",
    "                            f\"--nbuckets={NBUCKETS}\",\n",
    "                            f\"--lr={LR}\",\n",
    "                            f\"--nnsize={NNSIZE}\"\n",
    "                        ],\n",
    "                    },\n",
    "                }\n",
    "            ],\n",
    "        },\n",
    "    }\n",
    "    parent = f\"projects/{project}/locations/{location}\"\n",
    "    response = client.create_custom_job(parent=parent, custom_job=custom_job)\n",
    "    print(\"response:\", response)\n",
    "\n",
    "create_custom_job_sample(\n",
    "    project=PROJECT,\n",
    "    display_name=JOB_NAME,\n",
    "    image_uri=IMAGE_URI)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2022 Google Inc. Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-6.m96",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m96"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
